import cv2
import numpy as np
import torch
import torchvision.transforms as T
from PIL import Image
from torchvision.models.detection import fasterrcnn_resnet50_fpn
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor


class FasterRCNNDetector:
    def __init__(self, model_path, num_classes=3, confidence_threshold=0.5):
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.confidence_threshold = confidence_threshold

        # 加载模型
        self.model = fasterrcnn_resnet50_fpn()
        in_features = self.model.roi_heads.box_predictor.cls_score.in_features
        self.model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

        # 加载训练好的权重
        checkpoint = torch.load(model_path, map_location=self.device)
        self.model.load_state_dict(checkpoint['model_state_dict'])
        self.model.to(self.device)
        self.model.eval()

        # 类别映射
        self.classes = {0: 'background', 1: 'bolt', 2: 'nut'}

        # 图像转换
        self.transform = T.Compose([
            T.ToTensor(),
            T.Normalize(mean=[0.485, 0.456, 0.406],
                        std=[0.229, 0.224, 0.225])
        ])

    def detect(self, image_path):
        """
        检测图像中的目标
        """
        # 读取图像
        image = Image.open(image_path).convert('RGB')
        orig_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)

        # 转换图像
        image_tensor = self.transform(image)
        image_tensor = image_tensor.unsqueeze(0).to(self.device)

        with torch.no_grad():
            predictions = self.model(image_tensor)

        # 获取预测结果
        boxes = predictions[0]['boxes'].cpu().numpy()
        scores = predictions[0]['scores'].cpu().numpy()
        labels = predictions[0]['labels'].cpu().numpy()

        # 仅保留高置信度的预测
        mask = scores >= self.confidence_threshold
        boxes = boxes[mask]
        scores = scores[mask]
        labels = labels[mask]

        return orig_image, boxes, scores, labels

    def visualize(self, image, boxes, scores, labels):
        """
        可视化检测结果
        """
        # 颜色映射
        colors = {
            1: (0, 255, 0),  # bolt - 绿色
            2: (255, 0, 0)  # nut - 蓝色
        }

        # 复制图像以进行绘制
        image_draw = image.copy()

        # 绘制每个检测框
        for box, score, label in zip(boxes, scores, labels):
            # 获取坐标
            x1, y1, x2, y2 = map(int, box)

            # 绘制边界框
            color = colors.get(label, (0, 0, 255))
            cv2.rectangle(image_draw, (x1, y1), (x2, y2), color, 2)

            # 添加标签和置信度
            label_text = f'{self.classes[label]}: {score:.2f}'
            cv2.putText(image_draw, label_text, (x1, y1 - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

        return image_draw


def test_single_image(model_path, image_path, output_path=None, show=True):
    """
    测试单张图片
    """
    # 创建检测器
    detector = FasterRCNNDetector(model_path)

    # 执行检测
    image, boxes, scores, labels = detector.detect(image_path)

    # 可视化结果
    result_image = detector.visualize(image, boxes, scores, labels)

    # 保存结果
    if output_path:
        cv2.imwrite(output_path, result_image)

    # 显示结果
    if show:
        cv2.imshow('Detection Result', result_image)
        cv2.waitKey(0)
        cv2.destroyAllWindows()


def test_directory(model_path, image_dir, output_dir):
    """
    测试整个目录的图片
    """
    import os

    # 创建检测器
    detector = FasterRCNNDetector(model_path)

    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)

    # 获取所有图片
    image_files = [f for f in os.listdir(image_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]

    for image_file in image_files:
        # 构建完整路径
        image_path = os.path.join(image_dir, image_file)
        output_path = os.path.join(output_dir, f'detected_{image_file}')

        # 执行检测
        image, boxes, scores, labels = detector.detect(image_path)

        # 可视化结果
        result_image = detector.visualize(image, boxes, scores, labels)

        # 保存结果
        cv2.imwrite(output_path, result_image)
        print(f"Processed {image_file}")


if __name__ == '__main__':
    # 设置路径
    MODEL_PATH = '../pt_file/resnet_segmentation_final.pth'  # 修改为你的模型路径

    # 测试单张图片
    # IMAGE_PATH = 'path/to/your/test/image.jpg'  # 修改为你的测试图片路径
    # OUTPUT_PATH = 'output/detection_result.jpg'  # 修改为你想保存结果的路径
    #
    # test_single_image(MODEL_PATH, IMAGE_PATH, OUTPUT_PATH)

    # 或者测试整个目录

    IMAGE_DIR = 'lslm/lslm-test'  # 修改为你的测试图片目录
    OUTPUT_DIR = 'lslm/answer'  # 修改为你想保存结果的目录
    test_directory(MODEL_PATH, IMAGE_DIR, OUTPUT_DIR)
